import pandas as pd
import numpy as np

# 假设 data_handler.py 和 strategy.py 在同一个目录下
from data_handler import DataHandler
from strategy import Strategy

class BacktestEngine:
    """
    回测核心 (Backtesting Core) - V3 (适配ETF交易，记录更完善)
    """

    def __init__(self, start_date, end_date, data_handler, strategy, 
                 initial_capital=1000000, commission_rate=0,
                 risk_free_rate=0.04):
        """
        初始化回测引擎。

        :param start_date: 回测开始日期 (str)
        :param end_date: 回测结束日期 (str)
        :param data_handler: DataHandler 的实例 (需包含etf_close)
        :param strategy: Strategy 的实例 (适配新版线性策略)
        :param initial_capital: 初始资金 (float)
        :param commission_rate: 交易手续费率 (float)
        :param risk_free_rate: 年化无风险利率 (float)，用于计算夏普比率等
        """
        self.start_date = start_date
        self.end_date = end_date
        self.data_handler = data_handler
        self.strategy = strategy
        self.initial_capital = initial_capital
        self.commission_rate = commission_rate
        self.risk_free_rate = risk_free_rate

        # --- 账户状态变量 ---
        self.cash = initial_capital
        self.shares = 0
        self.total_assets = initial_capital
        
        # --- 记录器 ---
        self.history = []
        self.trade_records = []

    def run_backtest(self):
        """
        执行回测的主函数。
        """
        daily_data_generator = self.data_handler.get_data_generator(self.start_date, self.end_date)

        if not daily_data_generator:
            print("错误：指定日期范围内无数据，回测终止。")
            return None, None, None

        previous_day_data = None
        
        for today_data in daily_data_generator:
            # 【核心修改】所有资产估值和交易都基于ETF价格
            # 假设交易和估值都使用当日的收盘价，这是一个常见的简化
            current_price = today_data.etf_close
            
            # 1. 交易前，根据当日价格更新总资产
            current_total_assets = self.cash + self.shares * current_price
            
            # 2. 根据 T-1 日的数据生成 T 日的交易信号
            trade_signal = 0
            if previous_day_data:
                trade_signal = self.strategy.generate_trade_signal(previous_day_data)
            
            # 3. 执行交易
            self._execute_trade(trade_signal, current_total_assets, today_data)
            
            # 4. 记录当日收盘后的账户快照
            self._record_daily_snapshot(today_data, trade_signal)
            
            # 5. 更新前一日数据
            previous_day_data = today_data

        # 回测期结束后，清算所有头寸
        if previous_day_data:
            self._liquidate_position(previous_day_data)
        
        # 将历史记录转换为DataFrame
        history_df = pd.DataFrame(self.history)
        if not history_df.empty:
            history_df.set_index('date', inplace=True)
            
        # 计算详细的性能指标
        performance_metrics = self._calculate_performance_metrics(history_df)
        
        return performance_metrics, history_df, pd.DataFrame(self.trade_records)

    def _execute_trade(self, trade_signal, current_total_assets, today_data):
        if trade_signal == 0:
            return

        # 【核心修改】交易价格为当天的ETF收盘价
        trade_price = today_data.etf_close

        # 获取风控模块给出的最大允许仓位（基于指数点位）
        max_allowed_position = self.strategy.get_max_allowed_position(today_data)
        current_position_ratio = (self.shares * trade_price) / current_total_assets if current_total_assets > 0 else 0
        
        # 目标交易金额
        target_trade_value = current_total_assets * trade_signal

        if trade_signal > 0: # 买入逻辑
            # 如果目标仓位超过风控上限，则调整买入信号
            if current_position_ratio + trade_signal > max_allowed_position:
                adjusted_signal = max(0, max_allowed_position - current_position_ratio)
                target_trade_value = current_total_assets * adjusted_signal

            # 如果调整后的金额超过现金，则用尽所有现金
            if target_trade_value > self.cash:
                target_trade_value = self.cash
                
            if target_trade_value > 0:
                commission = target_trade_value * self.commission_rate
                actual_buy_value = target_trade_value - commission
                shares_to_buy = actual_buy_value / trade_price
                
                self.cash -= target_trade_value
                self.shares += shares_to_buy
                self._record_trade(today_data.Index, 'BUY', shares_to_buy, trade_price, commission)

        elif trade_signal < 0: # 卖出逻辑
            value_to_sell = abs(target_trade_value)
            shares_value_held = self.shares * trade_price

            # 不能卖出超过持仓的部分
            if value_to_sell > shares_value_held:
                value_to_sell = shares_value_held

            if value_to_sell > 0:
                shares_to_sell = value_to_sell / trade_price
                commission = value_to_sell * self.commission_rate

                self.cash += (value_to_sell - commission)
                self.shares -= shares_to_sell
                self._record_trade(today_data.Index, 'SELL', shares_to_sell, trade_price, commission)

    def _record_daily_snapshot(self, day_data, trade_signal):
        # 【核心修改】记录更完善的每日数据
        current_price = day_data.etf_close
        self.total_assets = self.cash + self.shares * current_price
        position_ratio = (self.shares * current_price) / self.total_assets if self.total_assets > 0 else 0
        
        self.history.append({
            'date': day_data.Index.to_pydatetime(),
            'index_close': day_data.close,       # 指数收盘价 (用于基准)
            'etf_close': current_price,          # ETF收盘价 (交易标的)
            'spread': day_data.spread,           # 添加股债利差
            'total_assets': self.total_assets,
            'cash': self.cash,
            'shares': self.shares,
            'position_ratio': position_ratio,
            'trade_signal': trade_signal         # 记录当日产生的信号
        })

    def _liquidate_position(self, last_day_data):
        if self.shares > 0:
            # 【核心修改】使用最后一天的ETF价格清算
            liquidation_price = last_day_data.etf_close
            sell_value = self.shares * liquidation_price
            commission = sell_value * self.commission_rate
            
            self.cash += (sell_value - commission)
            shares_sold = self.shares
            self.shares = 0
            self.total_assets = self.cash
            
            self._record_trade(last_day_data.Index, 'LIQUIDATE', shares_sold, liquidation_price, commission)
            
            # 更新历史记录的最后一天
            if self.history:
                self.history[-1].update({
                    'total_assets': self.total_assets, 'cash': self.cash,
                    'shares': 0, 'position_ratio': 0,
                })

    def _record_trade(self, date, side, shares, price, commission):
        self.trade_records.append({
            'date': date.to_pydatetime().date(), 'side': side, 'shares': shares,
            'price': price, 'commission': commission
        })

    def _calculate_performance_metrics(self, history_df):
        """
        【核心修改】计算更完善、更清晰的回测绩效指标。
        """
        if history_df.empty:
            return {}

        TRADING_DAYS_PER_YEAR = 252

        # 1. 收益率计算
        total_return = (history_df['total_assets'].iloc[-1] / self.initial_capital) - 1
        days = (history_df.index[-1] - history_df.index[0]).days
        years = max(days / 365.25, 1/TRADING_DAYS_PER_YEAR) # 避免回测期过短导致年化夸张
        
        # 【区分】复合年化收益率 (CAGR)
        annual_return_cagr = (1 + total_return) ** (1 / years) - 1
        
        history_df['daily_return'] = history_df['total_assets'].pct_change().fillna(0)
        
        # 【区分】算术平均年化收益率
        avg_daily_return = history_df['daily_return'].mean()
        annual_return_arithmetic = avg_daily_return * TRADING_DAYS_PER_YEAR

        # 2. 风险指标
        annual_volatility = history_df['daily_return'].std() * np.sqrt(TRADING_DAYS_PER_YEAR)
        
        # 3. 风险调整后收益指标
        daily_risk_free_rate = self.risk_free_rate / TRADING_DAYS_PER_YEAR
        excess_return = history_df['daily_return'] - daily_risk_free_rate
        
        sharpe_ratio = (excess_return.mean() / excess_return.std()) * np.sqrt(TRADING_DAYS_PER_YEAR) if excess_return.std() > 0 else 0

        # 4. 回撤分析
        history_df['peak'] = history_df['total_assets'].cummax()
        history_df['drawdown'] = (history_df['total_assets'] - history_df['peak']) / history_df['peak']
        max_drawdown = history_df['drawdown'].min()
        
        # 5. 其他常用指标
        calmar_ratio = annual_return_cagr / abs(max_drawdown) if max_drawdown < 0 else np.nan
        history_df['equity_curve'] = history_df['total_assets'] / self.initial_capital

        # 6. 基准表现 (使用指数)
        benchmark_start_price = history_df['index_close'].iloc[0]
        history_df['benchmark_equity_curve'] = history_df['index_close'] / benchmark_start_price
        benchmark_total_return = (history_df['index_close'].iloc[-1] / benchmark_start_price) - 1
        benchmark_annual_return = (1 + benchmark_total_return) ** (1 / years) - 1

        metrics = {
            'start_date': history_df.index[0].strftime('%Y-%m-%d'),
            'end_date': history_df.index[-1].strftime('%Y-%m-%d'),
            'total_return': total_return,
            'annualized_return_cagr': annual_return_cagr,
            'annualized_return_arithmetic': annual_return_arithmetic,
            'annual_volatility': annual_volatility,
            'sharpe_ratio': sharpe_ratio,
            'max_drawdown': max_drawdown,
            'calmar_ratio': calmar_ratio,
            'total_trades': len(self.trade_records),
            'benchmark_total_return': benchmark_total_return,
            'benchmark_annual_return': benchmark_annual_return,
        }
        return metrics